"""
基于Transformer的空调负荷预测工作流
支持可插拔的数据获取函数，便于不同数据源的切换
"""

import pandas as pd
import numpy as np
import datetime
import sys
import os
from typing import Callable, Dict, Tuple, Optional, Any
import warnings
warnings.filterwarnings('ignore')

# 导入增强版模块
from enhanced_hvac_preprocessing import EnhancedHVACDataProcessor
from enhanced_hvac_transformer import EnhancedHVACConfig, EnhancedHVACTransformer, RealTimeEnhancedPredictor
from enhanced_hvac_training import EnhancedHVACTrainer, load_enhanced_deployment_model


class HVACLoadPredictionWorkflow:
    """空调负荷预测工作流类"""
    
    def __init__(self, 
                 data_fetch_function: Callable = None,
                 model_config: EnhancedHVACConfig = None,
                 processor_config: dict = None):
        """
        初始化工作流
        
        Args:
            data_fetch_function: 数据获取函数，应该返回包含timestamp, swt, swh, real_load的DataFrame
            model_config: 模型配置
            processor_config: 数据处理器配置
        """
        # 设置默认数据获取函数
        self.data_fetch_function = data_fetch_function 
        
        # 设置默认模型配置
        self.model_config = model_config or EnhancedHVACConfig(
            d_model=256,
            n_heads=8,
            n_layers=6,
            seq_len=168,  # 一周历史数据
            pred_len=24,  # 预测24小时
            n_env_features=2,  # 温度和含湿量
            dropout=0.4
        )
        
        # 设置数据处理器配置
        self.processor_config = processor_config or {
            'seq_len': self.model_config.seq_len,
            'pred_len': self.model_config.pred_len,
            'outlier_threshold': 3.0,
            'missing_threshold': 0.3,
            'smooth_window': 3,
            'cyclical_encoding': True,
            'holiday_feature': True,
        }
        
        # 初始化组件
        self.processor = None
        self.model = None
        self.trainer = None
        self.predictor = None
        
        print(f"🚀 空调负荷预测工作流初始化完成")
        func_name = self.data_fetch_function.__name__ if self.data_fetch_function else "未设置"
        print(f"📊 数据获取函数: {func_name}")
        print(f"🏗️ 模型配置: {self.model_config.seq_len}h输入 → {self.model_config.pred_len}h预测")
    
    def set_data_fetch_function(self, func: Callable):
        """
        设置数据获取函数
        
        Args:
            func: 数据获取函数，应该接受start_time和end_time参数，
                  返回包含timestamp, swt, swh, real_load列的DataFrame
        """
        self.data_fetch_function = func
        func_name = func.__name__ if func else "None"
        print(f"📥 数据获取函数已更新为: {func_name}")
    
    def fetch_and_process_data(self, start_time: str, end_time: str) -> Dict[str, np.ndarray]:
        """
        获取并处理数据
        
        Args:
            start_time: 开始时间 "YYYY-MM-DD HH:MM:SS"
            end_time: 结束时间 "YYYY-MM-DD HH:MM:SS"
            
        Returns:
            processed_data: 处理后的数据字典
        """
        print(f"📅 获取数据: {start_time} 至 {end_time}")
        
        # 1. 调用数据获取函数
        if self.data_fetch_function is None:
            raise ValueError("未设置数据获取函数，请先调用 set_data_fetch_function() 或在初始化时提供")
        
        raw_data = self.data_fetch_function(start_time, end_time)
        
        if raw_data is None or len(raw_data) == 0:
            raise ValueError("数据获取函数返回空数据")
        
        print(f"📊 原始数据形状: {raw_data.shape}")
        print(f"📋 数据列名: {list(raw_data.columns)}")
        
        # 2. 数据格式检查和转换
        required_cols = ['timestamp', 'swt', 'swh', 'real_load']
        missing_cols = [col for col in required_cols if col not in raw_data.columns]
        if missing_cols:
            raise ValueError(f"数据缺少必要列: {missing_cols}")
        
        # 转换为标准格式
        data = raw_data[required_cols].copy()
        
        # 确保时间戳格式正确
        if not isinstance(data['timestamp'].iloc[0], pd.Timestamp):
            data['timestamp'] = pd.to_datetime(data['timestamp'])
        
        # 删除缺失值
        initial_len = len(data)
        data = data.dropna()
        if len(data) < initial_len:
            print(f"⚠️ 删除了 {initial_len - len(data)} 行缺失数据")
        
        if len(data) < self.model_config.seq_len + self.model_config.pred_len:
            raise ValueError(f"数据量不足，需要至少 {self.model_config.seq_len + self.model_config.pred_len} 个样本")
        
        # 3. 适配到模型格式
        # 负荷数据
        load_data = data['real_load'].values if hasattr(data['real_load'], 'values') else data['real_load']
        
        # 环境数据：将含湿量转换为相对湿度（简化处理）
        # 这里假设含湿量可以直接使用，实际应用中可能需要转换
        swt_data = data['swt'].values if hasattr(data['swt'], 'values') else data['swt']
        swh_data = data['swh'].values if hasattr(data['swh'], 'values') else data['swh']
        
        env_data = np.column_stack([
            swt_data,  # 温度
            swh_data * 1000  # 含湿量转换为合适的数值范围
        ])
        
        # 时间戳
        timestamp_data = data['timestamp'].values if hasattr(data['timestamp'], 'values') else data['timestamp']
        timestamps = pd.DatetimeIndex(timestamp_data)
        
        # 4. 初始化并拟合数据处理器
        if self.processor is None:
            self.processor = EnhancedHVACDataProcessor(self.processor_config)
            print("🔧 拟合数据预处理器...")
            self.processor.fit(load_data, env_data, timestamps)
        
        # 5. 处理数据
        processed_data = self.processor.transform(load_data, env_data, timestamps)
        
        print(f"✅ 数据处理完成")
        print(f"📈 负荷数据: {load_data.shape}, 范围: [{np.min(load_data):.2f}, {np.max(load_data):.2f}]")
        print(f"🌡️ 环境数据: {env_data.shape}")
        
        return processed_data, data
    
    def prepare_training_data(self, start_time: str, end_time: str) -> Tuple[Dict[str, np.ndarray], np.ndarray]:
        """
        准备训练数据
        
        Returns:
            X: 输入特征字典
            y: 目标序列
        """
        processed_data, _ = self.fetch_and_process_data(start_time, end_time)
        
        # 创建训练序列
        X, y = self.processor.create_sequences(processed_data)
        
        print(f"🎯 训练数据准备完成:")
        for key, value in X.items():
            print(f"  X_{key}: {value.shape}")
        print(f"  y: {y.shape}")
        
        return X, y
    
    def train_model(self, start_time: str, end_time: str, 
                   epochs: int = 100, batch_size: int = 32, 
                   learning_rate: float = 0.001, save_dir: str = "models") -> Dict:
        """
        训练模型
        
        Args:
            start_time: 训练数据开始时间
            end_time: 训练数据结束时间
            epochs: 训练轮数
            batch_size: 批次大小
            learning_rate: 学习率
            save_dir: 模型保存目录
            
        Returns:
            training_result: 训练结果
        """
        print(f"🚀 开始模型训练...")
        
        # 准备数据
        X, y = self.prepare_training_data(start_time, end_time)
        
        # 创建训练器
        self.trainer = EnhancedHVACTrainer(self.model_config, save_dir)
        
        # 训练模型
        training_result = self.trainer.train(
            X, y, 
            epochs=epochs, 
            batch_size=batch_size, 
            learning_rate=learning_rate
        )
        
        # 评估模型
        evaluation_result = self.trainer.evaluate(X, y, self.processor)
        
        print(f"📊 训练结果:")
        print(f"  最佳验证损失: {training_result['best_val_loss']:.6f}")
        print(f"  最终验证MAE: {training_result['final_val_mae']:.6f}")
        print(f"📈 模型性能:")
        for metric, value in evaluation_result['overall_metrics'].items():
            print(f"  {metric}: {value:.4f}")
        
        # 创建部署包
        package_dir = self.trainer.create_deployment_package(self.processor)
        training_result['package_dir'] = package_dir
        
        return training_result
    
    def load_trained_model(self, package_dir: str):
        """加载训练好的模型"""
        print(f"📦 加载训练好的模型: {package_dir}")
        
        self.predictor, self.processor = load_enhanced_deployment_model(package_dir)
        print(f"✅ 模型加载成功")
    
    def predict_with_realtime_data(self, historical_start: str, historical_end: str, 
                                 pred_hours: int = 24) -> Tuple[np.ndarray, Dict]:
        """
        使用实时数据进行预测
        
        Args:
            historical_start: 历史数据开始时间
            historical_end: 历史数据结束时间（当前时间）
            pred_hours: 预测小时数
            
        Returns:
            predictions: 预测结果
            stats: 预测统计信息
        """
        if self.predictor is None:
            raise ValueError("请先加载训练好的模型或完成模型训练")
        
        print(f"🔮 开始实时预测...")
        print(f"📅 历史数据: {historical_start} 至 {historical_end}")
        print(f"⏰ 预测未来 {pred_hours} 小时")
        
        # 获取历史数据
        _, raw_data = self.fetch_and_process_data(historical_start, historical_end)
        
        # 确保有足够的历史数据
        if len(raw_data) < self.model_config.seq_len:
            raise ValueError(f"历史数据不足，需要至少 {self.model_config.seq_len} 个时间点")
        
        # 使用最近的数据更新预测器
        recent_data = raw_data.tail(self.model_config.seq_len)
        
        for _, row in recent_data.iterrows():
            env_data = np.array([row['swt'], row['swh'] * 1000])
            self.predictor.update_data(
                load_value=row['real_load'],
                env_data=env_data,
                timestamp=row['timestamp']
            )
        
        # 进行预测
        if pred_hours == 24:
            # 预测今日剩余时间
            predictions, stats = self.predictor.predict_remaining_day()
        else:
            # 自定义预测长度（需要修改预测器以支持任意长度）
            predictions, stats = self.predictor.predict_remaining_day()
            if len(predictions) > pred_hours:
                predictions = predictions[:pred_hours]
                stats['remaining_hours'] = pred_hours
        
        print(f"✅ 预测完成")
        print(f"📊 预测统计:")
        print(f"  预测时长: {stats['remaining_hours']} 小时")
        print(f"  平均负荷: {stats['mean_load']:.2f}")
        print(f"  峰值负荷: {stats['max_load']:.2f}")
        print(f"  最低负荷: {stats['min_load']:.2f}")
        
        return predictions, stats
    
    def batch_predict(self, start_time: str, end_time: str, 
                     pred_hours: int = 24, step_hours: int = 1) -> pd.DataFrame:
        """
        批量预测
        
        Args:
            start_time: 开始时间
            end_time: 结束时间
            pred_hours: 每次预测的小时数
            step_hours: 预测步长（小时）
            
        Returns:
            results: 预测结果DataFrame
        """
        if self.predictor is None:
            raise ValueError("请先加载训练好的模型或完成模型训练")
        
        print(f"📊 开始批量预测...")
        
        # 生成预测时间点
        start_dt = pd.to_datetime(start_time)
        end_dt = pd.to_datetime(end_time)
        
        results = []
        current_time = start_dt
        
        while current_time < end_dt:
            try:
                # 计算历史数据时间范围
                hist_start = current_time - pd.Timedelta(hours=self.model_config.seq_len)
                hist_end = current_time
                
                # 进行预测
                predictions, stats = self.predict_with_realtime_data(
                    hist_start.strftime('%Y-%m-%d %H:%M:%S'),
                    hist_end.strftime('%Y-%m-%d %H:%M:%S'),
                    pred_hours
                )
                
                # 记录结果
                result = {
                    'prediction_time': current_time,
                    'pred_mean': stats['mean_load'],
                    'pred_max': stats['max_load'],
                    'pred_min': stats['min_load'],
                    'pred_hours': stats['remaining_hours']
                }
                
                # 添加每小时的预测值
                for i, pred_val in enumerate(predictions):
                    result[f'pred_hour_{i+1}'] = pred_val
                
                results.append(result)
                
                print(f"  ✅ {current_time}: 平均负荷 {stats['mean_load']:.2f}")
                
            except Exception as e:
                print(f"  ❌ {current_time}: 预测失败 - {e}")
            
            current_time += pd.Timedelta(hours=step_hours)
        
        return pd.DataFrame(results)


# 便捷函数
def create_workflow(data_fetch_function: Callable = None) -> HVACLoadPredictionWorkflow:
    """创建工作流实例"""
    return HVACLoadPredictionWorkflow(data_fetch_function)


def quick_train_and_predict(start_time: str, end_time: str, 
                          data_fetch_function: Callable = None,
                          epochs: int = 50) -> Dict:
    """
    快速训练和预测流程
    
    Args:
        start_time: 数据开始时间
        end_time: 数据结束时间
        data_fetch_function: 数据获取函数
        epochs: 训练轮数
        
    Returns:
        result: 包含训练结果和预测结果的字典
    """
    # 创建工作流
    workflow = create_workflow(data_fetch_function)
    
    # 计算数据分割点（80%用于训练）
    start_dt = pd.to_datetime(start_time)
    end_dt = pd.to_datetime(end_time)
    total_hours = (end_dt - start_dt).total_seconds() / 3600
    train_hours = int(total_hours * 0.8)
    
    train_end_dt = start_dt + pd.Timedelta(hours=train_hours)
    pred_start_dt = train_end_dt - pd.Timedelta(hours=workflow.model_config.seq_len)
    
    # 训练模型
    print(f"🎯 训练阶段: {start_time} 至 {train_end_dt}")
    training_result = workflow.train_model(
        start_time, 
        train_end_dt.strftime('%Y-%m-%d %H:%M:%S'),
        epochs=epochs
    )
    
    # 加载训练好的模型
    workflow.load_trained_model(training_result['package_dir'])
    
    # 进行预测
    print(f"🔮 预测阶段: 从 {train_end_dt} 开始预测")
    predictions, stats = workflow.predict_with_realtime_data(
        pred_start_dt.strftime('%Y-%m-%d %H:%M:%S'),
        train_end_dt.strftime('%Y-%m-%d %H:%M:%S')
    )
    
    return {
        'training_result': training_result,
        'predictions': predictions,
        'prediction_stats': stats,
        'workflow': workflow
    }


if __name__ == "__main__":
    # 示例使用
    print("🏢 空调负荷预测工作流演示")
    
    # 1. 使用默认数据获取函数
    print("\n=== 使用默认数据获取函数 ===")
    workflow = create_workflow()
    
    # 测试数据获取
    try:
        start_time = "2025-08-24 00:00:00"
        end_time = "2025-08-25 01:00:00"
        
        print("📊 测试数据获取...")
        processed_data, raw_data = workflow.fetch_and_process_data(start_time, end_time)
        print(f"✅ 数据获取成功，原始数据 {len(raw_data)} 条")
        
    except Exception as e:
        print(f"❌ 数据获取失败: {e}")
        print("💡 请确保数据库连接正常或使用模拟数据")
    
    # 2. 演示自定义数据获取函数
    print("\n=== 使用自定义数据获取函数 ===")
    
    def mock_data_fetch(start_time: str, end_time: str) -> pd.DataFrame:
        """模拟数据获取函数"""
        start_dt = pd.to_datetime(start_time)
        end_dt = pd.to_datetime(end_time)
        
        # 生成时间序列
        timestamps = pd.date_range(start_dt, end_dt, freq='H')
        n_points = len(timestamps)
        
        # 模拟负荷数据
        hours = np.arange(n_points) % 24
        base_load = 100 + 50 * np.sin(2 * np.pi * hours / 24 - np.pi/2)
        real_load = base_load + np.random.normal(0, 10, n_points)
        
        # 模拟环境数据
        swt = 25 + 10 * np.sin(2 * np.pi * np.arange(n_points) / (24*7)) + np.random.normal(0, 2, n_points)
        swh = 0.01 + 0.005 * np.random.rand(n_points)  # 含湿量
        
        return pd.DataFrame({
            'timestamp': timestamps,
            'swt': swt,
            'swh': swh,
            'real_load': real_load
        })
    
    # 使用自定义函数创建工作流
    mock_workflow = create_workflow(mock_data_fetch)
    
    # 快速训练和预测演示
    print("🚀 开始快速训练和预测演示...")
    
    try:
        # 生成一周的模拟数据进行演示
        demo_start = "2024-01-01 00:00:00"
        demo_end = "2024-01-08 00:00:00"
        
        result = quick_train_and_predict(
            demo_start, demo_end, 
            mock_data_fetch, 
            epochs=20  # 演示用较少epoch
        )
        
        print(f"\n✅ 演示完成!")
        print(f"📊 训练性能 - 验证MAE: {result['training_result']['final_val_mae']:.4f}")
        print(f"🔮 预测结果 - 平均负荷: {result['prediction_stats']['mean_load']:.2f}")
        
    except Exception as e:
        print(f"❌ 演示失败: {e}")
        print("💡 这可能是由于缺少依赖或计算资源不足")
    
    print("\n🎉 工作流演示结束!")
